
import os
import json
import tempfile
import argparse
from typing import Dict, Any

from magic_pdf.data.data_reader_writer import FileBasedDataReader, FileBasedDataWriter
from magic_pdf.data.dataset import PymuDocDataset
from magic_pdf.model.doc_analyze_by_custom_model import doc_analyze
from magic_pdf.config.enums import SupportedPdfParseMethod

from fastapi import FastAPI, UploadFile, File, HTTPException
from fastapi.responses import JSONResponse

# ------------------------------------------------------------
# 核心解析函数：给 bytes，返回 Python dict（中间 JSON）
# ------------------------------------------------------------
def parse_pdf_bytes(pdf_bytes: bytes) -> Dict[str, Any]:
    """
    1) 自动选 OCR/非 OCR
    2) 调用模型分析
    3) 返回中间 JSON（dict）
    """
    with tempfile.TemporaryDirectory() as tmpdir:
        # 保存临时 PDF
        pdf_path = os.path.join(tmpdir, "tmp.pdf")
        with open(pdf_path, "wb") as f:
            f.write(pdf_bytes)

        # 结果输出到同一目录
        out_dir = os.path.join(tmpdir, "out")
        os.makedirs(out_dir, exist_ok=True)

        # 读取 + 构建 Dataset
        reader = FileBasedDataReader("")             # 本地文件系统
        pdf_bytes2 = reader.read(pdf_path)
        ds = PymuDocDataset(pdf_bytes2)

        # 根据内容选模式
        if ds.classify() == SupportedPdfParseMethod.OCR:
            infer = ds.apply(doc_analyze, ocr=True)
            result = infer.pipe_ocr_mode(FileBasedDataWriter(out_dir))
        else:
            infer = ds.apply(doc_analyze, ocr=False)
            result = infer.pipe_txt_mode(FileBasedDataWriter(out_dir))

        # 读取中间 JSON
        json_path = os.path.join(out_dir, "result_middle.json")
        if not os.path.exists(json_path):
            raise RuntimeError("解析失败，未生成 JSON")
        with open(json_path, "r", encoding="utf-8") as jf:
            return json.load(jf)

# ------------------------------------------------------------
# 1) 本地脚本调用：CLI 接口
# ------------------------------------------------------------
def cli_main():
    parser = argparse.ArgumentParser(
        description="MinerU PDF 解析（命令行模式）"
    )
    parser.add_argument("pdf", help="待解析的 PDF 文件路径")
    parser.add_argument(
        "-o", "--out", help="输出 JSON 文件路径（不填则打印到 stdout）"
    )
    args = parser.parse_args()

    # 读取 PDF bytes 并解析
    with open(args.pdf, "rb") as f:
        pdf_bytes = f.read()
    try:
        result = parse_pdf_bytes(pdf_bytes)
    except Exception as e:
        print(f"Error: {e}", file=sys.stderr)
        exit(1)

    # 输出
    if args.out:
        with open(args.out, "w", encoding="utf-8") as wf:
            json.dump(result, wf, ensure_ascii=False, indent=2)
    else:
        print(json.dumps(result, ensure_ascii=False, indent=2))


# ------------------------------------------------------------
# 2) HTTP 接口调用：FastAPI
# ------------------------------------------------------------
app = FastAPI(title="MinerU PDF 解析服务")

@app.post("/parse_pdf")
async def parse_pdf_api(file: UploadFile = File(...)):
    if not file.filename.lower().endswith(".pdf"):
        raise HTTPException(400, "仅支持 PDF 格式")
    pdf_bytes = await file.read()
    try:
        result = parse_pdf_bytes(pdf_bytes)
    except Exception as e:
        raise HTTPException(500, f"解析失败: {e}")
    return JSONResponse(content=result)


# ------------------------------------------------------------
# 入口：根据是否有参数决定运行模式
# ------------------------------------------------------------
if __name__ == "__main__":
    import sys
    # 如果带了 --serve，就启动 API；否则当 CLI
    if "--serve" in sys.argv:
        # 向 uvicorn 传递 app
        # python mineru_service.py --serve --host 0.0.0.0 --port 8000
        import uvicorn
        parser = argparse.ArgumentParser()
        parser.add_argument("--host", default="0.0.0.0")
        parser.add_argument("--port", type=int, default=8000)
        # 其他 uvicorn 选项可自行扩展
        args, _ = parser.parse_known_args()
        uvicorn.run("mineru_service:app",
                    host=args.host, port=args.port, reload=False)
    else:
        cli_main()